#pragma once
#ifndef YOLOV5_H
#define YOLOV5_H
#include <iostream>
#include <opencv2/opencv.hpp>
#include <math.h>
using namespace std;
using namespace cv;
#include "json/json.h"

#define LIMIT_MIN(x) (x > 0 ? x : 0)
//模型的输出的结构
struct Output
{
    int id;           //结果类别id
    float confidence; //结果置信度
    cv::Rect box;     //矩形框

    Json::Value toJson(string name)
    {
        Json::Value jTmpValue;
        jTmpValue["x"] = int(LIMIT_MIN(box.x));
        jTmpValue["y"] = int(LIMIT_MIN(box.y));
        jTmpValue["width"] = int(box.width);
        jTmpValue["height"] = int(box.height);
        jTmpValue["name"] = name;
        jTmpValue["confidence"] = confidence;
        return jTmpValue;
    }
};
class datap
{
public:
    datap() {}
    std::vector<int> classIds;      //结果id数组
    std::vector<float> confidences; //结果每个id对应置信度数组
    std::vector<cv::Rect> boxes;    //每个id矩形框
    vector<int> nms_result;
    void reset();
    void press(float *pdata, size_t dlen = 61200, int net_width = 85, float conf_thres = 0.25, float nmsconf = 0.25, float *sc = 0);
    void print();
    void getbox(vector<Output> &output);
    // void getbox(vector<Outputf> &output);
};

class Yolov5
{
public:
    Yolov5();
    ~Yolov5(){};
    //设置Sigmoid函数
    float Sigmoid(float x)
    {
        return static_cast<float>(1.f / (1.f + exp(-x)));
    }
    datap dp;
    // anchors
    // stride
    bool isNeedAlert;

    const int netWidth = 1280; //网络模型输入大小
    const int netHeight = 1280;
    float nmsThreshold = 0.45;
    float boxThreshold = 0.25;
    float classThreshold = 0.15;
    float *pdata;             //存储模型推理的结果
    vector<cv::Scalar> color; //颜色列表
                              //    std::vector<std::string> className = { "dent", "scratch"};
    std::vector<std::string> className = {"fixed_stall", "sunshade", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
                                          "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
                                          "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
                                          "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
                                          "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
                                          "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
                                          "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
                                          "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
                                          "hair drier", "toothbrush"};
    string pack_json(std::vector<Output> &detectedout, bool jsonAlertCode);

    const int net_width = className.size() + 5;
    string mStrOutJson;
    string postresult_trt(float *pdata, cv::Mat img, vector<Output> &output, float thr, int wd); //处理TensorRT的结果
    cv::Rect get_rect(cv::Mat &img, float bbox[4]);
    //对结果画框并输出
    void drawPred(cv::Mat &img, std::vector<Output> result, std::vector<cv::Scalar> color);
};

#endif // YOLOV5_H
